GPU-accelerated SEACells (opt-in, end-to-end on GPU) - #83
Merged
Conversation
hussenmi
marked this pull request as ready for review
August 20, 2026 19:18
Adds a unified SEACellsModel that runs the metacell solver on the GPU (kernel build, Frank-Wolfe updates, and the objective) via two changes: - K residency: keep the kernel resident on the GPU instead of re-uploading it and evaluating the reconstruction error on the CPU. ~18x faster per iteration at 15k cells; cd34 end-to-end 110s -> 6.5s (~17x). - Reduced-form RSS: evaluate ||M - MBA|| via ||M||^2 - 2 tr(KBA) + tr(B^T K B A A^T), with K = M^T M, which needs O(n*s) memory instead of the dense n x n reconstruction. Fits on one GPU at 200k+ cells (~10 GB vs ~174 GB). Exact, not an approximation. Opt-in via core.SEACells(use_gpu=True, use_unified=True); use_unified defaults to False, so the existing CPU and legacy backends are unchanged. Adds parity tests (unified CPU == legacy cpu_dense; reduced-form RSS == direct Frobenius norm; GPU == CPU on a shared kernel). Same optimum: on cd34, CPU and GPU converge to the same RSS. Packaging via uv: `uv sync` (CPU) / `uv sync --extra gpu` (RAPIDS/CuPy/FAISS). uv.lock is gitignored rather than committed. See docs/gpu_speed_and_scale.md.
hussenmi
force-pushed
the
gpu-acceleration
branch
from
August 20, 2026 23:17
8d33ed5 to
1cc8c4c
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This adds a GPU version of SEACells. It's off by default; turn it on with
use_gpu=True, use_unified=Trueand nothing else changes:What's different on the GPU
The kernel stays on the GPU. The old GPU code kept the kernel on the CPU, copied it to the GPU every iteration, and computed the reconstruction error on the CPU. Keeping everything on the GPU makes each iteration about 18x faster. cd34 finishes in 6.5s instead of 110s.
The error avoids the big n×n matrix. The error used to be computed by building an n×n matrix, which gets huge at scale (174 GB for 200k cells). It's now computed with the identity
||M - MBA||^2 = ||M||^2 - 2 tr(KBA) + tr(B^T K B A A^T), which only needs n×s memory (about 10 GB for 200k cells). Same number, far less memory, so it runs on a single GPU at sizes the old code couldn't.Same results
The GPU path gives the same answer as the CPU path (see
tests/test_unified_parity.py), and on cd34 the CPU and GPU runs converge to the same error.Install
uv syncfor CPU, oruv sync --extra gpufor the GPU version (installs RAPIDS, CuPy, FAISS). Details indocs/gpu_speed_and_scale.md.